#!/usr/bin/env python3

import tkinter as tk
from tkinter import ttk
from tkinter import filedialog
import json
import rosbag
from BagPlayerProcess import BagPlayerProcess
import os
from ResultWinGUI import ResultWinTop
from TopicsSelectGUI import TopicsSelect
from CollectResultsGUI import CollectResults
import rospy


class ROSBagPlayerGUI:
    def __init__(self) -> None:
        self.rosbag_process = None
        # 创建主窗口
        self.root = tk.Tk()
        self.root.title("ROS Bag Player")

        self.json_file = os.path.abspath(__file__)
        self.json_file = os.path.join(
            os.path.dirname(self.json_file), "config", "default.json"
        )
        self.configuration = self.readconfig(self.json_file)

        self.playying_bag = None
        self.last_bag_statue = "STOPPED"
        default_bag = self.configuration["last_bag"]
        self.bag_total_time = 0

        self.param_updata_tick = 0
        self.params = {}
        if not os.path.exists(default_bag):
            default_bag = ""
        else:
            bag = rosbag.Bag(default_bag)
            _, self.bag_total_time = self.get_bag_info(bag)
        self.menu_bar = tk.Menu(self.root)
        self.file_menu = tk.Menu(self.menu_bar, tearoff=0)
        self.file_menu.add_command(label="Open(O)", command=self.browse_folder)
        self.file_menu.add_command(label="Open Last(P)", command=self.open_last_bag)
        self.file_menu.add_command(label="Generate PDF", command=self.export_results)
        self.file_menu.add_separator()
        self.menu_bar.add_cascade(label="File", menu=self.file_menu)
        self.root.config(menu=self.menu_bar)

        # 创建标签和文本输入框
        self.path_label = tk.Label(self.root, text="ROS Bag Path:", font=("Ariel", 10))
        self.path_label.grid(row=0, column=0, padx=5, pady=5)

        self.bagpath_label = tk.Label(
            self.root,
            font=("Ariel", 8),
            text=default_bag,
            wraplength=500,
            justify="left",
        )
        self.bagpath_label.grid(row=1, column=0, columnspan=4, padx=5, pady=5)

        self.topic_select_button = tk.Button(
            self.root, text="Select Topics", command=self.select_publish_topics
        )
        self.topic_select_button.grid(row=0, column=4, rowspan=2, padx=5, pady=5)

        # 创建播放按钮
        self.play_button = tk.Button(self.root, text="play", command=self.play_rosbag)
        self.play_button.grid(row=2, column=0, rowspan=2, padx=5, pady=5)

        self.pause_button = tk.Button(
            self.root, text="pause", command=self.pause_rosbag, state="disabled"
        )
        self.pause_button.grid(row=2, column=1, rowspan=2, padx=5, pady=5)

        self.progress_bar = ttk.Progressbar(
            self.root, orient="horizontal", length=200, mode="determinate"
        )
        self.progress_bar.grid(
            row=2, column=2, columnspan=2, padx=5, pady=5, sticky="s"
        )

        self.bag_info_text = tk.StringVar(value="STOPPED")
        self.bag_info_label = tk.Label(self.root, textvariable=self.bag_info_text)
        self.bag_info_label.grid(
            row=3, column=2, columnspan=2, rowspan=2, padx=5, pady=5, sticky="n"
        )

        self.rate_label = tk.Label(self.root, text="Publish rate:")
        self.rate_label.grid(row=2, column=4, rowspan=1, padx=5, pady=5, sticky="s")

        self.rate_spinbox = tk.Spinbox(
            self.root, width=5, from_=0.1, to=3, increment=0.1
        )
        self.rate_spinbox.grid(row=3, column=4, padx=5, pady=5, sticky="n")

        self.command_label = tk.Label(self.root, text="")
        self.command_label.grid(
            row=4, column=0, padx=5, pady=5, columnspan=6, sticky="w"
        )

        self.root.bind("<space>", self.pause_rosbag)
        self.root.bind("<o>", lambda x: self.browse_folder())
        self.root.bind("<p>", lambda x: self.open_last_bag())
        self.root.bind("O", lambda x: self.browse_folder())
        self.root.bind("P", lambda x: self.open_last_bag())
        self.selected_topics = self.configuration["selected_topics"]
        self.publish_rate = self.configuration["publish_rate"]

        self.rate_spinbox.config(textvariable=tk.StringVar(value=self.publish_rate))

        self.bag_player_process = BagPlayerProcess()
        self.update_bag_statue()

    def open_last_bag(self):
        last_bag = self.configuration["last_bag"]
        if not os.path.exists(last_bag):
            return
        bag_file = self.find_latest_bagfile(last_bag)
        if bag_file:
            self.bagpath_label.config(text=bag_file, wraplength=500, justify="left")
            bag = rosbag.Bag(bag_file)
            topics, self.bag_total_time = self.get_bag_info(bag)
            self.select_topics(topics)
            self.rate_spinbox.config(state="normal")

    def find_latest_bagfile(self, file_path):
        folder_path = os.path.dirname(file_path)
        files_with_extension = [
            f for f in os.listdir(folder_path) if f.endswith(".bag")
        ]

        if not files_with_extension:
            print("No files with the specified extension found.")
            return None

        files_with_extension_full_paths = [
            os.path.join(folder_path, f) for f in files_with_extension
        ]

        sorted_files = sorted(
            files_with_extension_full_paths,
            key=lambda x: os.path.getmtime(x),
            reverse=True,
        )

        return sorted_files[0]

    def readconfig(self, config_path):
        try:
            with open(config_path, "r") as file:
                data = json.load(file)
        except:
            self.CreateConfigJSON(config_path)
            with open(config_path, "r") as file:
                data = json.load(file)
        return data

    def CreateConfigJSON(self, config_path):
        base_path = os.path.dirname(config_path)
        if not os.path.exists(base_path):
            os.makedirs(base_path)
        configuration_default = {
            "last_bag": "",
            "selected_topics": [],
            "publish_rate": 1.0,
        }
        with open(config_path, "w") as file:
            json.dump(configuration_default, file, indent=4)

    def saveconfig(self, config_path):
        self.configuration["last_bag"] = self.bagpath_label.cget("text")
        self.configuration["selected_topics"] = self.selected_topics
        self.configuration["publish_rate"] = self.rate_spinbox.get()
        with open(config_path, "w") as file:
            json.dump(self.configuration, file, indent=4)

    def select_publish_topics(self):
        bag_path = self.bagpath_label.cget("text")
        if not bag_path.endswith(".bag"):
            print("Please input the rosbag file path.")
            return
        bag = rosbag.Bag(bag_path)
        topics, self.bag_total_time = self.get_bag_info(bag)
        self.select_topics(topics)

    def get_bag_info(self, bag: rosbag.Bag):
        topics = bag.get_type_and_topic_info().topics.keys()
        self.bag_total_time = bag.get_end_time() - bag.get_start_time()
        return topics, self.bag_total_time

    def browse_folder(self):
        initialfile = self.configuration["last_bag"]
        initdir = os.path.dirname(initialfile)
        folder_path = filedialog.askopenfilename(
            initialfile=initialfile, initialdir=initdir, filetypes=[("rosbag", "*.bag")]
        )  # 弹出文件夹选择对话框
        if len(folder_path) == 0 or not folder_path.endswith(".bag"):
            return
        if folder_path:
            self.bagpath_label.config(text=folder_path, wraplength=500, justify="left")
            bag = rosbag.Bag(folder_path)
            # 获取bag文件中的所有topic
            topics, self.bag_total_time = self.get_bag_info(bag)
            self.rate_spinbox.config(state="normal")
            self.select_topics(topics)

    def select_topics(self, topics):
        topics_window = TopicsSelect(
            self.root, topics, self.selected_topics, self.topic_select_callback
        )
        topics_window.grab_set()

    def topic_select_callback(self, selected_topics):
        self.selected_topics = selected_topics

    def play_rosbag(self):
        rosbag_file_path = self.bagpath_label.cget("text")
        self.playying_bag = rosbag_file_path
        if not rosbag_file_path or not rosbag_file_path.endswith(".bag"):
            print("Please input the rosbag file path.")
            return
        if self.bag_player_process.is_running:
            print("Bag player is already running.")
            return
        self.saveconfig(self.json_file)
        self.play_button.config(text="stop", command=self.stop_rosbag)
        self.pause_button.config(state="active")
        self.rate_spinbox.config(state="disabled")
        publish_rate = str(self.rate_spinbox.get())
        command = self.bag_player_process.SetBagInfo(
            rosbag_file_path, self.selected_topics, publish_rate
        )
        self.command_label.config(
            text=f"Command:\n {command}",
            font=("Arial", 8),
            wraplength=700,
            justify="left",
        )
        self.bag_player_process.start()

    def stop_rosbag(self, result_win_top=True):
        if self.bag_player_process.is_running:
            self.bag_player_process.stop()
            self.bag_player_process.join()
            self.progress_bar["value"] = 0
            self.pause_button.config(state="disabled")
            self.rate_spinbox.config(state="normal")
            self.play_button.config(text="play", command=self.play_rosbag)
            self.bag_player_process = BagPlayerProcess()
            if result_win_top:
                result_win = ResultWinTop(self.root, self.playying_bag, self.params)
                result_win.grab_set()

    def pause_rosbag(self, event=None):
        self.bag_player_process.set_pause(True)

    def export_results(self):
        bag_path = self.bagpath_label.cget("text")
        if not bag_path or not bag_path.endswith(".bag"):
            print("Please input the rosbag file path.")
            return
        result_win = CollectResults(self.root, bag_path)
        result_win.grab_set()

    def set_bag_process(self, bag_time, bag_cur_time):
        progress_value = 0
        bag_info_label = ""
        if self.bag_total_time == 0:
            progress_value = 0
            bag_info_label = f"--: --/--"
        else:
            progress_value = float(bag_cur_time) / float(self.bag_total_time) * 100
            bag_info_label = f"{float(bag_time):.2f}: {float(bag_cur_time):.2f}/{float(self.bag_total_time):.2f}"
        self.progress_bar["value"] = progress_value
        self.bag_info_label.config(
            textvariable=tk.StringVar(value=bag_info_label), font=("Arial", 8)
        )

    def update_rosparam(self):
        self.param_updata_tick += 1
        if self.param_updata_tick >= 10:
            self.param_updata_tick = 0
            params = rospy.get_param_names()
            for param in params:
                if not param in self.params.keys():
                    self.params[param] = rospy.get_param(param)

    def update_bag_statue(self):
        is_running, bag_time, bag_cur_time, bag_total_time = (
            self.bag_player_process.GetStatue()
        )
        self.set_bag_process(bag_time, bag_cur_time)
        if is_running == "STOPPED" and not self.last_bag_statue == "STOPPED":
            self.stop_rosbag()
        elif is_running == "FAILED_ROSCORE":
            tk.messagebox.showerror(
                "Error", "Failed to start roscore, please start roscore first."
            )
            self.stop_rosbag(False)
        elif is_running == "RUNNING":
            self.update_rosparam()
        self.root.after(100, self.update_bag_statue)
        self.last_bag_statue = is_running

    def run(self):
        self.root.mainloop()

    def stop(self):
        if self.bag_player_process.is_running:
            self.bag_player_process.stop()
            self.bag_player_process.join()
        self.root.quit()
